package com.pine.dbant.common.plugin.meta.function.jdbc;

import com.google.common.collect.Lists;
import com.pine.dbant.common.model.DatabaseTermTransformParams;
import com.pine.dbant.common.model.JdbcMetadataTreeView;
import com.pine.dbant.common.model.PrimaryKey;
import com.pine.dbant.common.plugin.meta.JdbcMetaAnalysisException;
import com.pine.dbant.common.plugin.meta.JdbcMetadataAnalysisConfig;
import com.pine.dbant.common.plugin.meta.QueryMetadataFunction;
import com.pine.dbant.common.plugin.meta.QueryMetadataParam;
import lombok.extern.slf4j.Slf4j;

import java.sql.DatabaseMetaData;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.List;

/**
 * @author xiaoyuan
 * @create 2023/3/17 11:25
 **/
@Slf4j
public class PrimaryKeyFunction implements QueryMetadataFunction<JdbcMetadataAnalysisConfig, List> {

    @Override
    public  List<PrimaryKey> query(JdbcMetadataAnalysisConfig metadataAnalysisConfig, QueryMetadataParam queryMetadataParam) {
        DatabaseMetaData dbMetaData;
        try {
            dbMetaData = metadataAnalysisConfig.getConnection().getMetaData();
        } catch (SQLException e) {
            throw new JdbcMetaAnalysisException(e.getMessage(), e);
        }
        String catalog = queryMetadataParam.getConfiguration().getNecessaryValue("catalog");
        String tableName = queryMetadataParam.getConfiguration().getNecessaryValue("tableName");
        List<PrimaryKey> primaryKeys = Lists.newArrayList();
        DatabaseTermTransformParams queryParams = transformQueryCatalogAndSchema(catalog, tableName);
        try (ResultSet rs = dbMetaData.getPrimaryKeys(queryParams.getCatalog(), queryParams.getSchema(), tableName)) {
            while (rs.next()) {
                primaryKeys.add(new PrimaryKey(rs.getString("COLUMN_NAME"), rs.getShort("KEY_SEQ"), rs.getString("PK_NAME")));
            }
        } catch (SQLException e) {
            log.error("获取数据库[{}]表[{}]的主键信息错误", catalog, tableName, e);
        }
        return primaryKeys;
    }

    @Override
    public String getType() {
        return JdbcMetadataTreeView.PRIMARY_KEY.getType();
    }
}
